import struct
from typing import List, Dict, Any, Optional
from datetime import datetime


class BluetoothProtocol:
    """
    蓝牙通信协议处理类
    基于README.md中描述的协议格式：55AA + 长度 + ID + 命令 + 数据 + CRC8
    """
    
    # 协议常量
    FRAME_HEADER = 0x55AA
    SERVICE_UUID = "FEE0"
    WRITE_CHARACTERISTIC = "FEE2"
    NOTIFY_CHARACTERISTIC = "FEE4"
    
    # 命令定义
    CMD_HANDSHAKE = 0x11
    CMD_STATUS_QUERY = 0x12
    CMD_HEARTBEAT = 0x13
    CMD_SET_MODE = 0x22
    CMD_SET_INTENSITY = 0x23
    CMD_SET_DURATION = 0x24
    CMD_TREATMENT_CONTROL = 0x41
    
    # 治疗模式
    MODE_TDCS_POSITIVE = 0x01  # tDCS+
    MODE_TDCS_NEGATIVE = 0x02  # tDCS-
    MODE_TACS = 0x03           # tACS
    
    # 治疗控制状态
    TREATMENT_START = 0x01
    TREATMENT_PAUSE = 0x02
    TREATMENT_STOP = 0x03
    TREATMENT_RESUME = 0x04
    
    def __init__(self):
        self.frame_id = 0x01
    
    def _calculate_crc8(self, data: bytes) -> int:
        """计算CRC8校验码"""
        crc = 0x00
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x80:
                    crc = (crc << 1) ^ 0x07
                else:
                    crc <<= 1
                crc &= 0xFF
        return crc
    
    def _build_frame(self, command: int, data: bytes = b'') -> bytes:
        """构建数据帧"""
        # 帧头 (2字节) + 长度 (1字节) + ID (1字节) + 命令 (1字节) + 数据 + CRC8 (1字节)
        length = 4 + len(data)  # ID + 命令 + 数据长度
        
        frame = struct.pack('>H', self.FRAME_HEADER)  # 帧头，大端序
        frame += struct.pack('B', length)             # 长度
        frame += struct.pack('B', self.frame_id)      # 帧ID
        frame += struct.pack('B', command)            # 命令
        frame += data                                 # 数据
        
        # 计算CRC8（不包括帧头）
        crc_data = frame[2:]  # 从长度字段开始
        crc = self._calculate_crc8(crc_data)
        frame += struct.pack('B', crc)
        
        # 更新帧ID
        self.frame_id = (self.frame_id + 1) % 256
        
        return frame
    
    def _parse_frame(self, frame: bytes) -> Optional[Dict[str, Any]]:
        """解析接收到的数据帧"""
        if len(frame) < 6:  # 最小帧长度
            return None
        
        try:
            # 解析帧头
            header = struct.unpack('>H', frame[:2])[0]
            if header != self.FRAME_HEADER:
                return None
            
            # 解析其他字段
            length = frame[2]
            frame_id = frame[3]
            command = frame[4]
            
            # 检查帧长度
            if len(frame) != length + 3:  # 长度 + 帧头(2) + 长度字段(1)
                return None
            
            # 提取数据和CRC
            data = frame[5:-1] if length > 2 else b''
            received_crc = frame[-1]
            
            # 验证CRC
            crc_data = frame[2:-1]
            calculated_crc = self._calculate_crc8(crc_data)
            if received_crc != calculated_crc:
                return None
            
            return {
                'frame_id': frame_id,
                'command': command,
                'data': data,
                'timestamp': datetime.utcnow()
            }
            
        except (struct.error, IndexError):
            return None
    
    def handshake(self) -> bytes:
        """握手命令"""
        handshake_data = struct.pack('>I', 0x11223344)  # 握手数据：11 22 33 44
        return self._build_frame(self.CMD_HANDSHAKE, handshake_data)
    
    def query_status(self) -> bytes:
        """查询设备状态"""
        return self._build_frame(self.CMD_STATUS_QUERY, b'\x01')
    
    def heartbeat(self, runtime: int = 0) -> bytes:
        """心跳命令"""
        heartbeat_data = struct.pack('>I', runtime)  # 运行时间（秒）
        return self._build_frame(self.CMD_HEARTBEAT, heartbeat_data)
    
    def set_treatment_mode(self, mode: str) -> bytes:
        """设置治疗模式"""
        mode_map = {
            'tDCS+': self.MODE_TDCS_POSITIVE,
            'tDCS-': self.MODE_TDCS_NEGATIVE,
            'tACS': self.MODE_TACS
        }
        
        if mode not in mode_map:
            raise ValueError(f"不支持的治疗模式: {mode}")
        
        mode_data = struct.pack('B', mode_map[mode])
        return self._build_frame(self.CMD_SET_MODE, mode_data)
    
    def set_current_intensity(self, intensity: float) -> bytes:
        """设置电流强度（0-2.0mA，精度0.1mA）"""
        if not 0 <= intensity <= 2.0:
            raise ValueError("电流强度必须在0-2.0mA范围内")
        
        # 转换为0-99的档位（0.1mA精度）
        level = int(intensity * 10)
        if level > 99:
            level = 99
        
        intensity_data = struct.pack('B', level)
        return self._build_frame(self.CMD_SET_INTENSITY, intensity_data)
    
    def set_treatment_duration(self, duration_minutes: int) -> bytes:
        """设置治疗时间（分钟）"""
        if not 1 <= duration_minutes <= 60:
            raise ValueError("治疗时间必须在1-60分钟范围内")
        
        # 转换为秒数，高字节在前
        duration_seconds = duration_minutes * 60
        duration_data = struct.pack('>H', duration_seconds)
        return self._build_frame(self.CMD_SET_DURATION, duration_data)
    
    def set_frequency(self, frequency: float) -> bytes:
        """设置频率（仅tACS模式，0.1-100Hz）"""
        if not 0.1 <= frequency <= 100.0:
            raise ValueError("频率必须在0.1-100Hz范围内")
        
        # 频率数据，使用2字节表示，精度0.1Hz
        freq_value = int(frequency * 10)
        freq_data = struct.pack('>H', freq_value)
        return self._build_frame(self.CMD_SET_MODE, freq_data)  # 复用模式设置命令
    
    def start_treatment(self) -> bytes:
        """开始治疗"""
        control_data = struct.pack('BB', 0x01, self.TREATMENT_START)  # 工作方式 + 状态
        return self._build_frame(self.CMD_TREATMENT_CONTROL, control_data)
    
    def pause_treatment(self) -> bytes:
        """暂停治疗"""
        control_data = struct.pack('BB', 0x01, self.TREATMENT_PAUSE)
        return self._build_frame(self.CMD_TREATMENT_CONTROL, control_data)
    
    def resume_treatment(self) -> bytes:
        """恢复治疗"""
        control_data = struct.pack('BB', 0x01, self.TREATMENT_RESUME)
        return self._build_frame(self.CMD_TREATMENT_CONTROL, control_data)
    
    def stop_treatment(self) -> bytes:
        """停止治疗"""
        control_data = struct.pack('BB', 0x01, self.TREATMENT_STOP)
        return self._build_frame(self.CMD_TREATMENT_CONTROL, control_data)
    
    def parse_device_response(self, response_data: bytes) -> Dict[str, Any]:
        """解析设备响应数据"""
        parsed = self._parse_frame(response_data)
        if not parsed:
            return {'error': '无效的响应数据'}
        
        command = parsed['command']
        data = parsed['data']
        result = {
            'command': command,
            'timestamp': parsed['timestamp'],
            'frame_id': parsed['frame_id']
        }
        
        try:
            if command == self.CMD_HANDSHAKE:
                # 握手响应
                if len(data) >= 4:
                    response_code = struct.unpack('>I', data[:4])[0]
                    result['handshake_success'] = response_code == 0x44332211
                
            elif command == self.CMD_STATUS_QUERY:
                # 状态查询响应
                if len(data) >= 8:
                    status_data = struct.unpack('>BBBBHBB', data[:8])
                    result.update({
                        'device_status': status_data[0],
                        'treatment_mode': status_data[1],
                        'current_intensity': status_data[2] / 10.0,  # 转换为mA
                        'battery_level': status_data[3],
                        'remaining_time': status_data[4],  # 秒
                        'treatment_status': status_data[5],
                        'error_code': status_data[6]
                    })
                
            elif command == self.CMD_HEARTBEAT:
                # 心跳响应
                if len(data) >= 4:
                    device_time = struct.unpack('>I', data[:4])[0]
                    result['device_runtime'] = device_time
                
            elif command == self.CMD_TREATMENT_CONTROL:
                # 治疗控制响应
                if len(data) >= 2:
                    control_result = struct.unpack('BB', data[:2])
                    result.update({
                        'control_result': control_result[0],
                        'treatment_status': control_result[1]
                    })
            
            else:
                # 其他命令的通用响应
                result['raw_data'] = data.hex()
                
        except (struct.error, IndexError) as e:
            result['parse_error'] = str(e)
            result['raw_data'] = data.hex()
        
        return result
    
    def create_treatment_session(self, params: Dict[str, Any]) -> List[bytes]:
        """创建完整的治疗会话命令序列"""
        commands = []
        
        # 1. 握手
        commands.append(self.handshake())
        
        # 2. 设置治疗模式
        if 'mode' in params:
            commands.append(self.set_treatment_mode(params['mode']))
        
        # 3. 设置电流强度
        if 'intensity' in params:
            commands.append(self.set_current_intensity(params['intensity']))
        
        # 4. 设置治疗时间
        if 'duration' in params:
            commands.append(self.set_treatment_duration(params['duration']))
        
        # 5. 设置频率（仅tACS模式）
        if params.get('mode') == 'tACS' and 'frequency' in params:
            commands.append(self.set_frequency(params['frequency']))
        
        # 6. 开始治疗
        commands.append(self.start_treatment())
        
        return commands
    
    def validate_treatment_params(self, params: Dict[str, Any]) -> Optional[str]:
        """验证治疗参数"""
        # 检查必需参数
        required_params = ['mode', 'intensity', 'duration']
        for param in required_params:
            if param not in params:
                return f"缺少必需参数: {param}"
        
        # 验证治疗模式
        valid_modes = ['tDCS+', 'tDCS-', 'tACS']
        if params['mode'] not in valid_modes:
            return f"无效的治疗模式: {params['mode']}"
        
        # 验证电流强度
        try:
            intensity = float(params['intensity'])
            if not 0 <= intensity <= 2.0:
                return "电流强度必须在0-2.0mA范围内"
        except (ValueError, TypeError):
            return "电流强度必须是数字"
        
        # 验证治疗时间
        try:
            duration = int(params['duration'])
            valid_durations = [10, 15, 20, 30]
            if duration not in valid_durations:
                return f"治疗时间必须是 {valid_durations} 中的一个"
        except (ValueError, TypeError):
            return "治疗时间必须是整数"
        
        # 验证频率（tACS模式）
        if params['mode'] == 'tACS':
            if 'frequency' not in params:
                return "tACS模式需要设置频率"
            try:
                frequency = float(params['frequency'])
                if not 0.1 <= frequency <= 100.0:
                    return "频率必须在0.1-100Hz范围内"
            except (ValueError, TypeError):
                return "频率必须是数字"
        
        return None  # 验证通过


# 使用示例
if __name__ == "__main__":
    protocol = BluetoothProtocol()
    
    # 创建治疗参数
    treatment_params = {
        'mode': 'tDCS+',
        'intensity': 1.5,
        'duration': 20
    }
    
    # 验证参数
    validation_error = protocol.validate_treatment_params(treatment_params)
    if validation_error:
        print(f"参数验证失败: {validation_error}")
    else:
        # 创建命令序列
        commands = protocol.create_treatment_session(treatment_params)
        print(f"生成了 {len(commands)} 个命令")
        
        # 打印命令的十六进制表示
        for i, cmd in enumerate(commands):
            print(f"命令 {i+1}: {cmd.hex().upper()}")